package sub

import (
  "bufio"
  "bytes"
  "encoding/binary"
  "fmt"
  "os"
  "strconv"
  "strings"
  "unsafe"
  "time"

  bpf "github.com/iovisor/gobpf/bcc"
)

import "C"

type EventType int32

const (
  eventArg EventType = iota
  eventRet
)

const source string = `
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>
#include <linux/fs.h>

#define ARGSIZE  128

enum event_type {
    EVENT_ARG,
    EVENT_RET,
};

struct data_t {
    u32 pid;  // PID as in the userspace term (i.e. task->tgid in kernel)
    u32 ppid; // Parent PID as in the userspace term (i.e task->real_parent->tgid in kernel)
    u32 uid;
    char comm[TASK_COMM_LEN];
    enum event_type type;
    char argv[ARGSIZE];
    int retval;
};

BPF_PERF_OUTPUT(events);

static int __submit_arg(struct pt_regs *ctx, void *ptr, struct data_t *data)
{
    bpf_probe_read_user(data->argv, sizeof(data->argv), ptr);
    events.perf_submit(ctx, data, sizeof(struct data_t));
    return 1;
}

static int submit_arg(struct pt_regs *ctx, void *ptr, struct data_t *data)
{
    const char *argp = NULL;
    bpf_probe_read_user(&argp, sizeof(argp), ptr);
    if (argp) {
        return __submit_arg(ctx, (void *)(argp), data);
    }
    return 0;
}

int syscall__execve(struct pt_regs *ctx,
    const char __user *filename,
    const char __user *const __user *__argv,
    const char __user *const __user *__envp)
{

    u32 uid = bpf_get_current_uid_gid() & 0xffffffff;

    // if (container_should_be_filtered()) {
    //     return 0;
    // }

    // create data here and pass to submit_arg to save stack space (#555)
    struct data_t data = {};
    struct task_struct *task;

    data.pid = bpf_get_current_pid_tgid() >> 32;

    task = (struct task_struct *)bpf_get_current_task();
    // Some kernels, like Ubuntu 4.13.0-generic, return 0
    // as the real_parent->tgid.
    // We use the get_ppid function as a fallback in those cases. (#1883)
    data.ppid = task->real_parent->tgid;

    bpf_get_current_comm(&data.comm, sizeof(data.comm));
    data.type = EVENT_ARG;

    __submit_arg(ctx, (void *)filename, &data);

    // skip first arg, as we submitted filename
    #pragma unroll
    for (int i = 1; i < MAX_ARGS; i++) {
        if (submit_arg(ctx, (void *)&__argv[i], &data) == 0)
             goto out;
    }

    // handle truncated argument list
    char ellipsis[] = "...";
    __submit_arg(ctx, (void *)ellipsis, &data);
out:
    return 0;
}

int do_ret_sys_execve(struct pt_regs *ctx)
{
    // if (container_should_be_filtered()) {
    //     return 0;
    // }

    struct data_t data = {};
    struct task_struct *task;

    u32 uid = bpf_get_current_uid_gid() & 0xffffffff;

    data.pid = bpf_get_current_pid_tgid() >> 32;
    data.uid = uid;

    task = (struct task_struct *)bpf_get_current_task();
    // Some kernels, like Ubuntu 4.13.0-generic, return 0
    // as the real_parent->tgid.
    // We use the get_ppid function as a fallback in those cases. (#1883)
    data.ppid = task->real_parent->tgid;

    bpf_get_current_comm(&data.comm, sizeof(data.comm));
    data.type = EVENT_RET;
    data.retval = PT_REGS_RC(ctx);
    events.perf_submit(ctx, &data, sizeof(data));

    return 0;
}`

type execveEvent struct {
  Pid    uint32
  Ppid   uint32
  Uid   uint32
  Comm   [16]byte
  Type   int32
  Argv   [128]byte
  RetVal int32
}

type OneoffProcessInfo struct {
	UnixMilli   int64 `json:"time,omitempty"`
	Comm   string `json:"comm"`
	Pid    uint32 `json:"pid"`
	Ppid   uint32 `json:"ppid"`
	Uid    uint32 `json:"uid"`
	Argv   string `json:"argv"`
	RetVal int32  `json:"retval"`
}

// getPpid is a fallback to read the parent PID from /proc.
// Some kernel versions, like 4.13.0 return 0 getting the parent PID
// from the current task, so we need to use this fallback to have
// the parent PID in any kernel.
func getPpid(pid uint32) uint32 {

  f, err := os.OpenFile(fmt.Sprintf("/proc/%d/status", pid), os.O_RDONLY, os.ModePerm)
  if err != nil {
    return 0
  }
  defer f.Close()

  sc := bufio.NewScanner(f)
  for sc.Scan() {
    text := sc.Text()
    if strings.Contains(text, "PPid:") {
      f := strings.Fields(text)
      i, _ := strconv.ParseUint(f[len(f)-1], 10, 32)
      return uint32(i)
    }
  }
  return 0

}

func Main0() {
  // 1. make output channel
  ear := make(chan OneoffProcessInfo)
  stopButton := make(chan struct{})

  // 2. pass it
  go ListenToOneOff(ear, stopButton)
  fmt.Println("Press 'Enter' to exit...")
  scn := bufio.NewScanner(os.Stdin)

  go func() {
    // the consumer 🦜: Now it's a printer
    for {
      fmt.Println(<-ear)
    }
  }()

  scn.Scan()
  fmt.Println("👋 bye.")
  stopButton <- struct{}{}
}

func ListenToOneOff(out chan<- OneoffProcessInfo, stopButton <-chan struct{}){

  // --------------------------------------------------
  // 1. boilerplate
  m := bpf.NewModule(strings.Replace(source, "MAX_ARGS", strconv.FormatUint(20, 10), -1), []string{})
  //MaxArgs: 20

  defer m.Close()

  fnName := bpf.GetSyscallFnName("execve")

  kprobe, err := m.LoadKprobe("syscall__execve")
  if err != nil {
    fmt.Fprintf(os.Stderr, "Failed to load syscall__execve: %s\n", err)
    os.Exit(1)
  }

  // passing -1 for maxActive signifies to use the default
  // according to the kernel kprobes documentation
  if err := m.AttachKprobe(fnName, kprobe, -1); err != nil {
    fmt.Fprintf(os.Stderr, "Failed to attach syscall__execve: %s\n", err)
    os.Exit(1)
  }

  kretprobe, err := m.LoadKprobe("do_ret_sys_execve")
  if err != nil {
    fmt.Fprintf(os.Stderr, "Failed to load do_ret_sys_execve: %s\n", err)
    os.Exit(1)
  }

  // passing -1 for maxActive signifies to use the default
  // according to the kernel kretprobes documentation
  if err := m.AttachKretprobe(fnName, kretprobe, -1); err != nil {
    fmt.Fprintf(os.Stderr, "Failed to attach do_ret_sys_execve: %s\n", err)
    os.Exit(1)
  }

  table := bpf.NewTable(m.TableId("events"), m)

  channel := make(chan []byte, 1000)

  perfMap, err := bpf.InitPerfMap(table, channel, nil)
  // 🦜 Here we add the channel to kernal (the BPF vm to be more specific)

  if err != nil {
    fmt.Fprintf(os.Stderr, "Failed to init perf map: %s\n", err)
    os.Exit(1)
  }

  go func() {

    args := make(map[uint32][]string)

    for {
      data := <-channel // 🦜 Here we receive the data from kernal

      var event execveEvent
      err := binary.Read(bytes.NewBuffer(data), bpf.GetHostByteOrder(), &event)

      if err != nil {
        fmt.Printf("failed to decode received data: %s\n", err)
        continue
      }

      if eventArg == EventType(event.Type) {
        e, ok := args[event.Pid]
        if !ok {
          e = make([]string, 0)
        }
        argv := (*C.char)(unsafe.Pointer(&event.Argv))

        e = append(e, C.GoString(argv))
        args[event.Pid] = e
      } else {
        if event.RetVal != 0 && !true {
          delete(args, event.Pid)
          continue
        }

        comm := C.GoString((*C.char)(unsafe.Pointer(&event.Comm)))
        argv, ok := args[event.Pid]
        if !ok {
          continue
        }

        p := OneoffProcessInfo{
          // YYYY-MM-DD HH:MM:SS
          UnixMilli:   time.Now().UnixMilli(),
          Pid:    event.Pid,
          Ppid:   event.Ppid,
          Uid:    event.Uid,
          Comm:   comm,
          RetVal: event.RetVal,
        }

        if event.Ppid == 0 {
          event.Ppid = getPpid(event.Pid)
        }

        p.Argv = strings.Join(argv, " ")
        p.Argv = strings.TrimSpace(strings.Replace(p.Argv, "\n", "\\n", -1))

        // fmt.Println(p)
        out <- p
        delete(args, event.Pid)
      }
    }
  }()

  perfMap.Start()
  <-stopButton
  perfMap.Stop()
}
